# -*- coding: utf-8 -*-
import datetime
import time
import pandas as pd
import os
import logging
try:
    import queue
except ImportError:
    import Queue as queue

from .data import TushareDataHandler, MongodbDataHandler
from .error import ParamException
from .logger import get_logger, set_logger
from .plotter import plotly
from .event import EventHub, EVENT, Event

class Backtest(object):
    """
    initial_capital: 初始资金
    execution_handler: 执行处理的类
    portfolio: 投资组合
    strategy: 自己实现的策略的类
    commission_type: 手续费函数
    slippage_type: 滑点
    **param: 其它参数, 为的是给自己提供的类提供参数
    """
    def __init__(self, initial_capital, execution_handler, portfolio, strategy,
                 commission_type='zero', slippage_type='zero', **params):

        # context就是自己, 目前可以通过这个获取整个系统所有的参数.
        self.context = self

        #策略开始日期和结束日期
        if 'end_date' in params:
            self.end_date = params['end_date']
        else:
            self.end_date = datetime.datetime.now()

        if 'start_date' in params:
            self.start_date = params['start_date']
        else:
            self.start_date = datetime.datetime.now() - datetime.timedelta(days=90)

        # 策略的初始化资金.(注意: 这个东西只用于初始化data-handler类, 并不会再变化, 不要把它当成当前的现金!)
        self.cash = float(initial_capital)

        # 事件引擎
        self.event_hub = EventHub()
        # 股票列表. 只有在stocks里面的股票才能够得到数据的更新!
        self.stocks = []

        # 各种数据的处理类赋值
        self.execution_handler_cls = execution_handler
        self.portfolio_cls = portfolio
        self.strategy_cls = strategy

        # 交易费和滑点的赋值
        self.commission_type = commission_type
        self.slippage_type = slippage_type

        # 初始化策略(因为策略的运行需要事件引擎)
        self.params = params
        self.strategy = self.strategy_cls(self.event_hub, self.context)
        self.existing_error = False

    def setup_logger(self, log_level=logging.INFO, to_file=""):
        # 设置自己的logger, 支持让logger输出到文件
        logger = logging.getLogger(__name__)
        logger.setLevel(log_level)
        formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')

        if to_file != "":
            file_handler = logging.FileHandler(os.path.join(os.getcwd(), to_file), mode='w')
            file_handler.setLevel(log_level)
            file_handler.setFormatter(formatter)
            logger.addHandler(file_handler)

        stream_handler = logging.StreamHandler()
        stream_handler.setLevel(log_level)
        stream_handler.setFormatter(formatter)
        logger.addHandler(stream_handler)
        set_logger(logger)

    def event_generator(self):
        bars = self.data_handler
        while True:
            # update time
            if not bars.continue_backtest:
                break
            # 更新参照股票, 主要是为了更新时间
            bars.update_ref()
            self.now = bars.now
            if bars.continue_backtest:
                # 产生事件 before_bar
                yield Event(EVENT.BEFORE_BAR, context=self.context, now=self.now)
                # 获取当天的股票数据
                bar_dict = bars.update_bars()
                # 产生事件bar
                yield Event(EVENT.BAR, bar_dict=bar_dict, context=self.context, now=self.now)
                # 产生事件after_bar
                yield Event(EVENT.AFTER_BAR, bar_dict=bar_dict, context=self.context, now=self.now)
                get_logger().info("current cash: " + str(self.portfolio.current_holdings['cash']))


    def _run(self):
        # 主事件循环
        for event in self.event_generator():
            # 当得到事件之后, 执行对应的处理函数
            self.event_hub.publish_event(event)


    def _force_clear(self, event):
        # 强制平仓
        get_logger().debug("=====do force clear=====")
        for s in self.stocks:
            amount = self.portfolio.current_positions[s]
            print(self.data_handler.check_stop(s))
            self.strategy.order_shares(s, -amount)

        get_logger().debug("=====force clear end=====")

    def _output_performance(self):
        pass

    def _check_environment(self):
        # check 变量有没有初始化全
        # 变量有没有问题

        if self.start_date >= self.end_date:
            raise ParamException("start_date must before end_date")

        if self.stocks is None or len(self.stocks) == 0:
            raise ParamException("there is no stock in stocks")

        if self.cash is None or self.cash < 0:
            raise ParamException("cash mush > 0")

    def _build_instance(self):
        # 初始化一些数据处理类
        self.data_handler = MongodbDataHandler(self.event_hub, self.stocks, self.start_date, self.end_date, self.params)
        # self.data_handler = TushareDataHandler(self.event_hub, self.stocks, self.start_date, self.end_date, self.params)
        self.portfolio = self.portfolio_cls(self.event_hub, self.data_handler, self.start_date, self.cash)
        self.execution_handler = self.execution_handler_cls(self.event_hub, self.data_handler,
                                                            slippage_type=self.slippage_type,
                                                            commission_type=self.commission_type)

    def _register_listeners(self):
        # 注册事件及其相应的事件处理函数
        self.event_hub.add_listener(EVENT.BEFORE_BAR, self.strategy.before_bar_event)
        self.event_hub.add_listener(EVENT.BAR, self.strategy.handle_bar_event)
        self.event_hub.add_listener(EVENT.AFTER_BAR, self.strategy.after_bar_event)
        self.event_hub.add_listener(EVENT.AFTER_BAR, self.portfolio.update_timeindex)

        self.event_hub.add_listener(EVENT.PENDING_ORDER, self.portfolio.update_signal)
        self.event_hub.add_listener(EVENT.RAW_ORDER, self.execution_handler.execute_order)
        self.event_hub.add_listener(EVENT.ACCEPT_ORDER, self.portfolio.update_fill)

        self.event_hub.add_listener(EVENT.END_BACKTEST, self._force_clear)
        self.event_hub.add_listener(EVENT.END_BACKTEST, self.portfolio.update_timeindex)



    def _prepare(self):
        # 回测所需要做的准备工作
        try:
            # 策略自定义的init函数, 设置一些自定义的变量到context中
            self.strategy.init(self.context)
            # check 所有必须的变量有没有赋值完全
            self._check_environment()
            # 初始化数据处理类
            self._build_instance()
            # 注册事件和事件处理函数
            self._register_listeners()
        except Exception as e:
            print(e)
            self.existing_error = True

    # 回测的启动函数
    def run_backtest(self):
        # prepare, 初始化各种数据处理类. 初始化事件引擎. check有没有必要的数据没有进行赋值.如果没有的话报error
        self._prepare()
        if self.existing_error:
            return

        start_time = time.time()
        # 运行回测
        self._run()
        end_time = time.time()
        # 回测结束过后, 为了便于统计, 强制对股票进行平仓. 卖出所有股票
        self.event_hub.publish_event(Event(EVENT.END_BACKTEST))
        timing = round(end_time - start_time, 2)

        # log end
        # 提供的展示数据的接口
        self._output_performance()

        #不应该作为参数返回
        positions = pd.DataFrame(self.portfolio.all_positions).drop_duplicates(subset='datetime', keep='last'
                                                                               ).set_index('datetime')
        holdings = pd.DataFrame(self.portfolio.all_holdings).drop_duplicates(subset='datetime', keep='last'
                                                                             ).set_index('datetime')
        return positions, holdings

    def set_cash(self, cash):
        # 设置初始化资金
        self.cash = cash

    def set_duration(self, start_date, end_date):
        # 设置开始日期和结束日期. 类型为datetime.datetime()
        self.start_date = start_date
        self.end_date = end_date

    def plot(self, engine='plotly'):
        # 对数据进行画图
        go = plotly(self.context)
        go.plot()

